module Zeno.Solver
(
  solver
)
where

import Prelude ()
import StdImports
import qualified Data.Set as Set

import Zeno.Core
import Zeno.Type
import Zeno.Term
import Zeno.Clause
import Zeno.Function
import Zeno.ProofLogic ( Proof (..) )
import qualified Zeno.ProofLogic as PL

{- 
TODO

Do things in phases?
  like after induction there is no generalisation?
  after splitting there is no induction?
  
Remember to add kterm hypotheses to the conditions
  for elem-min
  
Theory about removal of multiple linked induction hypotheses on application

Theory about only adding a hypothesis one the "first" inductive step of a chain
  that is to say do not add a hypothesis if hypotheses already exist
  
Try merging case-splitting back into induction
  just remember to apply function definitions whilst finding critical terms
  
Try iding each case-analysis and never use case-analysis to do induction on
  a child of a variable that was created with this analysis

-}

data SolverState
  = SolverState   { sZeno :: !Zeno,
                    sGoal :: !ZClause,
                    sHypotheses :: ![ZClause],
                    sDoGeneralisation :: !Bool,
                    sIndSteps :: !Int,
                    sConSteps :: !Int,
                    sRandom :: !StdGen }
    
type Solver = Reader SolverState 
type StatefulSolver = State SolverState
type SolverStep = Solver [ZProof]

liftZenoReader :: ZenoReader a -> Solver a
liftZenoReader zr = do
  zeno <- sZeno <$> ask
  return $ runReader zr zeno
  
liftZenoState :: ZenoState a -> StatefulSolver a
liftZenoState zs = do
  zeno <- sZeno <$> get
  let (a, zeno') = runState zs zeno
  modify $ \s -> s { sZeno = zeno' }
  return a
  
solver :: (Int, Int) -> String -> StdGen -> ZenoReader ZProof
solver (ind, con) lemma_name random = do
  func <- lookupFunctionName lemma_name
  func_s <- showZFunction func
  let lemmas = clausifyFunction func
  when (null lemmas) $ error $ "Lemma " ++ lemma_name ++ " not found."
  lemmas_s <- intercalate "\n" <$> mapM showZClause lemmas  
  proofs <- traceZ ("solving " ++ func_s ++ "\nlemmas:\n" ++ lemmas_s) $ mapM solver' lemmas
  return $ case proofs of
    [] -> Unknown
    [proof] -> proof
    proofs -> PL.conjoinStep (PL.ProofStep "proofs" Nothing []) proofs
  where
    solver' :: ZClause -> ZenoReader ZProof
    solver' lemma = do
      zeno <- ask
      meep <- showZClause lemma
      let state = traceZ ("solving " ++ meep) $
            SolverState { sZeno = zeno,
                          sGoal = lemma { clauseVariables = [] },
                          sHypotheses = mempty,
                          sDoGeneralisation = True,
                          sIndSteps = ind,
                          sConSteps = con,
                          sRandom = random }
      let proof = runReader (addStep "goal" [] solve) state
      return $ if PL.proven proof then traceZ "Proven!" proof else proof
    
solve :: Solver ZProof
solve = do
  eq <- clauseEquality <$> sGoal <$> ask
  conds <- reducingConds
  let eq' = mapPair (substituteWithin (Set.fromList conds)) eq 
  (goal', changed) <- liftZenoReader $ 
    applyTermDefinitions conds (Clause [] eq' conds)
    
  let steps = concat <$> randomSequence
                [solve_hgn, solve_hyp, solve_fac, 
                 solve_gen, solve_cse, solve_ind]
      proof = checkGoal $ PL.anyKnown <$> steps
  if not changed
    then proof
    else localGoal goal' $ addStep "def" [] $ proof
  where
    reducingConds :: Solver [ZEquality]
    reducingConds = do
      conds <- clauseConditions <$> sGoal <$> ask
      let (red_r, rest) = partition (constructorTerm . snd) conds
          (red_l, _) = partition (constructorTerm . fst) rest
      return $ red_r ++ (map flipPair red_l)

solve_hyp :: Solver [ZProof]
solve_hyp = enableGeneralisation $ cannotDisprove $ do
  applyHypotheses maybeApplyHyp
  where
    maybeApplyHyp :: ZClause -> Solver [ZProof]
    maybeApplyHyp hyp = do
      goal@(l_goal, r_goal) <- clauseEquality <$> sGoal <$> ask
      let apply_left = findApplications l_goal hyp
          apply_right = findApplications r_goal (flipClauseEquality hyp)
          l_eqs = map (\t -> (t, r_goal)) apply_left
          r_eqs = map (\t -> (l_goal, t)) apply_right
          all_eqs = l_eqs ++ r_eqs
      hyp_s <- liftZenoReader $ showZClause hyp
      let step = "hyp " ++ hyp_s
      mapM (\eq -> localGoalEquality eq $ addStep step [] solve) all_eqs
    
    findApplications :: ZTerm -> ZClause -> [ZTerm]
    findApplications goal_term 
                     (Clause vars (from_hyp, to_hyp) []) = 
      map applyUnifier $ 
        filter (completeUnifier vars) $ 
          allUnifiers from_hyp goal_term
      where 
        applyUnifier :: ZTermSubstitution -> ZTerm
        applyUnifier uni = 
          let hyp_uni = mapPair (substituteWithin uni) (from_hyp, to_hyp)
          in substituteWithin (Set.singleton hyp_uni) goal_term
          
          
solve_fac :: Solver [ZProof]
solve_fac = cannotDisprove $ do
  goal <- sGoal <$> ask
  let (l_term,  r_term) = clauseEquality goal
      (l_func : l_args) = flattenTerm l_term
      (r_func : r_args) = flattenTerm r_term
  if l_func /= r_func
    then return [Unknown]
    else do
      step <- proofstep "fct" []
      let factor_goals = l_args `zip` r_args
          factor_proofs = mapM 
            (\eq -> localGoalEquality eq $ addStep "fac" [] solve) factor_goals
      return <$> PL.allProvenStep step <$> factor_proofs
      
      
solve_hgn :: Solver [ZProof]
solve_hgn = applyHypotheses generaliseHyp
  where
    generaliseHyp :: ZClause -> Solver [ZProof]
    generaliseHyp (Clause vars eq@(l_hyp, r_hyp) []) 
      | not (constructorTerm r_hyp) = return [Unknown]
      | otherwise = do
          goal_term <- fst <$> clauseEquality <$> sGoal <$> ask 
          
          let possibleTerm :: ZTerm -> Bool
              possibleTerm term = 
                not (isVar term) &&
                functionTerm term &&
                destructibleTerm term &&
                length (variables term) == length vars + 1
    
          sequence $ do
            possible <- filter possibleTerm (strictlyWithinList l_hyp)
            uni <- allUnifiers possible goal_term
            if not (completeUnifier vars uni) 
              then return (return Unknown)
              else do
                let eq' = mapPair (substituteWithin uni) eq
                    genr = substituteWithin uni possible
                return $ addCond eq' $ 
                         addStep "hgn" [] $ 
                         applyGeneralisation genr
      

solve_gen :: Solver [ZProof]
solve_gen = generalisation $ do
  inds <- filter variableTerm <$> criticalTerms
  goal <- sGoal <$> ask
  let superterms = concatMap (duplicateSuperterms (withinList goal)) inds
  mapM applyGeneralisation superterms
  where  
    duplicateSuperterms :: [ZTerm] -> ZTerm -> [ZTerm]
    duplicateSuperterms all_terms sub_term = 
      removeFactors all_duplicates
      where
        all_duplicates :: [ZTerm]
        all_duplicates = filter superterm (duplicates all_terms)
        
        superterm :: ZTerm -> Bool
        superterm t = destructibleTerm t && isApp t 
          && t `contains` sub_term && functionTerm t
        
          
solve_cse :: Solver [ZProof]
solve_cse = disableGeneralisation $ doIndStep $ cannotDisprove $ do
  splits <- filter isApp <$> criticalTerms
  mapM goSplit splits
  where
    goSplit :: ZTerm -> Solver ZProof
    goSplit term 
      | not (destructibleTerm term) = return Unknown
      | otherwise = do
          splits <- proveOverConstructors goSplit' (termType term)
          term_s <- liftZenoReader $ showZTerm term
          step <- proofstep "cse" []
          return $ PL.allProvenStep step splits
      where
        goSplit' :: (ZTerm, [ZVar]) -> Solver ZProof
        goSplit' (kterm, _) = do
          Clause [] eq conds <- sGoal <$> ask
          let goal = Clause [] eq ((term, kterm) : conds)
          
          term_s <- liftZenoReader $ showZTerm term
          kterm_s <- liftZenoReader $ showZTerm kterm
          let step = "cse " ++ term_s ++ " => " ++ kterm_s
          localGoal goal $ 
            takeIndStep $ addStep step [] solve


solve_ind :: Solver [ZProof]
solve_ind = disableGeneralisation $ doIndStep $ do
  inds <- map fromVar . filter variableTerm <$> criticalTerms
  mapM goInductive inds
  where
    goInductive :: ZVar -> Solver ZProof
    goInductive var 
      | not (destructibleVar var) = return Unknown
      | otherwise = do
          inds <- proveOverConstructors goInductive' (zVarType var)
          var_s <- liftZenoReader $ showZVar var
          step <- proofstep "ind" []
          return $ PL.conjoinStep step inds
      where
        goInductive' :: (ZTerm, [ZVar]) -> Solver ZProof
        goInductive' (ind, ind_vars) = do
          goal <- sGoal <$> ask
          let ind_sub = Set.singleton (Var var, ind)
          
              neededCondition :: ZEquality -> Bool
              neededCondition (t1, t2) = 
                var `elem` toList t1 || var `elem` toList t2
              
              hyp_goal = goal 
                { clauseConditions = 
                    filter neededCondition (clauseConditions goal) }
              
              createHyp :: ZVar -> StatefulSolver ZClause
              createHyp ind_var = liftZenoState $ do
                let others = filter (/= var) (variables hyp_goal)
                others' <- mapM freshVariable others
                let sublist = (var, ind_var) : (others `zip` others')
                    sub = Set.fromList $ map (Var *** Var) sublist
                    hyp = substituteWithin sub hyp_goal
                return $ hyp { clauseVariables = others' }
                
          var_s <- liftZenoReader $ showZVar var
          ind_s <- liftZenoReader $ showZTerm ind
          let step = "ind " ++ var_s ++ " => " ++ ind_s
                
          stateful (mapM createHyp ind_vars) $ \hyps -> do
            substituteTerms ind_sub $ 
              addHypotheses hyps $ 
                takeIndStep $
                  addStep step hyps solve

proveOverConstructors :: ((ZTerm, [ZVar]) -> Solver a) -> ZType -> Solver [a] 
proveOverConstructors prove ztype = do
  let ((TypeVar dtype@(ZDataType _ dt_args cons)) : ty_args) = flattenAppType ztype
      bindings = Set.fromList $ (map TypeVar dt_args) `zip` ty_args  
  stateful (mapM (liftZenoState . instantiate bindings) cons) (mapM prove)
  
applyGeneralisation :: ZTerm -> Solver ZProof
applyGeneralisation term = do
  stateful (liftZenoState newIdS) $ \new_id -> do
    let new_var = ZVar new_id (termType term) ZVariable
        term_sub = Set.singleton (term, Var new_var)
    term_s <- liftZenoReader $ showZTerm term
    var_s <- liftZenoReader $ showZVar new_var
    let step = "gen " ++ term_s ++ " => " ++ var_s
    substituteTerms term_sub $
      addStep step [] solve
              
completeUnifier :: [ZVar] -> ZTermSubstitution -> Bool
completeUnifier vars uni =
  Set.fromList (map Var vars) == Set.mapMonotonic fst uni
              
criticalTerms :: Solver [ZTerm]
criticalTerms = do
  Clause [] (l_term, r_term) conds <- sGoal <$> ask
  let all_terms = l_term : r_term : (map fst conds)
  filter destructibleTerm . nub <$> mapMaybeM criticalTerm all_terms
  where
    criticalTerm :: ZTerm -> Solver (Maybe ZTerm)
    criticalTerm term = do
      maybe_expr <- liftZenoReader $ applyFunction [] term
      case maybe_expr of
        Nothing -> 
          return (Just term)
        Just (Value {}) -> 
          return Nothing
        Just (CaseOf term' _) -> do
          if term `contains` term'
            then criticalTerm term'
            else return (Just term')

-- Applicable hypotheses are those with all conditions fulfilled
--   the given hypotheses should be conditionless
applyHypotheses :: (ZClause -> Solver [ZProof]) -> Solver [ZProof]
applyHypotheses apply = do
  hyps <- sHypotheses <$> ask
  concat <$> concatMapM applyHyp hyps
  where
    applyHyp :: ZClause -> Solver [[ZProof]]
    applyHyp cls = do
      apps <- applicable cls
      removeHypothesis cls (mapM apply apps)
  
    applicable :: ZClause -> Solver [ZClause]
    applicable cls@(Clause _ _ []) = return [cls]
    applicable (Clause c_vars c_eq c_conds) = do
      unis <- mapM removeCondition c_conds
      let unis' = map Set.unions (sequence unis)
      return $ do
        uni <- unis'
        let vars = map (fromVar . fst) (toList uni)
            c_vars' = filter (not . flip elem vars) c_vars
            c_eq' = mapPair (substituteWithin uni) c_eq
        return (Clause c_vars' c_eq' [])
      where
        removeCondition :: ZEquality -> Solver [ZTermSubstitution]
        removeCondition (c_l, c_r) = do
          g_conds <- clauseConditions <$> sGoal <$> ask
          return $ do
            (g_l, g_r) <- g_conds
            let uni_l = unify c_l g_l
                uni_r = unify c_r g_r
                uni = uni_l `Set.union` uni_r
            if isOneToMany uni && validUnifier uni
              then [uni]
              else []
          where 
            validUnifier :: ZTermSubstitution -> Bool
            validUnifier = all (flip elem (map Var c_vars) . fst)
          
fullGoal :: Solver ZClause
fullGoal = sGoal <$> ask
            
addStep :: String -> [ZClause] -> Solver ZProof -> Solver ZProof
addStep name added solver = do
  goal <- fullGoal
  PL.addStep (PL.ProofStep name (Just goal) added) <$> solver
  
cannotDisprove :: Functor f => f [ZProof] -> f [ZProof]
cannotDisprove = fmap (filter (not . PL.disproven))
            
-- Check whether the goal is trivially true or false
checkGoal :: Solver ZProof -> Solver ZProof
checkGoal solver = do
  goal@(Clause [] (t1, t2) conds) <- sGoal <$> ask
  steps <- sConSteps <$> ask
  meep <- liftZenoReader $ showZClause goal
  contr2 <- contradiction2 conds steps
  contr1 <- contradiction1 conds steps
  -- Check whether the conditions are contradictory
  if contr2
      then return $ Proven (PL.ProofStep "con" Nothing []) []
    else if PL.proven contr1
      then return contr1
    else if equalTerms t1 t2
      then return $ Proven (PL.ProofStep "eql" Nothing []) []
    else if unequalTerms t1 t2
      then return $ Disproven (PL.ProofStep "irf" Nothing []) []
    else do
      -- Filter out trivial conditions
      let conds' = filter (uncurry (/=)) conds
      localGoal (Clause [] (t1, t2) conds') solver
            
contradiction1 :: [ZEquality] -> Int -> Solver ZProof
contradiction1 eqs depth 
  | depth <= 0 = return Unknown
  | otherwise = do
      inverses <- mapM solveInverse (pickEach eqs)
      return (PL.anyProven inverses)
  where 
    solveInverse :: (ZEquality, [ZEquality]) -> Solver ZProof
    solveInverse (eq@(t1, t2), others)
      | not (constructorTerm t2) = return Unknown
      | otherwise = do
        let ((Var con) : _) = flattenTerm t2
            ((TypeVar dtype) : _) = flattenAppType (termType t2)
            conterms = map Var $ filter (/= con) (zDataTypeCons dtype)
            conterms' = filter varTypeTerm conterms
        PL.anyProven <$> mapM solveSimply conterms'
      where
        solveSimply :: ZTerm -> Solver ZProof
        solveSimply t2' = do
          let goal = Clause [] (t1, t2') others
          proof <- localGoal goal (simpleSolver solve)
          return $ PL.addStep (PL.ProofStep "con" (Just goal) []) proof
          
    simpleSolver :: Solver a -> Solver a
    simpleSolver = local $ \s -> s 
      { sConSteps = 0,
        sIndSteps = depth,
        sHypotheses = mempty,
        sDoGeneralisation = True }

-- Splitting based contradition finder
--   probably subsumed by the new solver based one
contradiction2 :: [ZEquality] -> Int -> Solver Bool
contradiction2 eqs depth = do
  let (terms, kterms) = unzip eqs
  terms' <- liftZenoReader $ 
    map fst <$> mapM (applyTermDefinitions []) terms
  let eqs' = (terms' `zip` kterms)
  if any (uncurry unequalTerms) eqs' 
      then return True
    else if all (uncurry equalTerms) eqs' || depth <= 0
      then return False
    else do
      term_splits <- concatMapM splits (filter varTypeTerm terms')
      results <- mapM (splitOn eqs') term_splits
      return (or results)
  where
    splitOn :: [ZEquality] -> ZTerm -> Solver Bool
    splitOn eqs term
      | not (destructibleTerm term) = return False
      | otherwise =
          and <$> proveOverConstructors testSplit (termType term)
      where
        testSplit :: (ZTerm, [ZVar]) -> Solver Bool
        testSplit (term', _) = do
          let sub = substituteWithin (Set.singleton (term, term'))
          contradiction2 (map (first sub) eqs) (depth - 1)
    
    splits :: ZTerm -> Solver [ZTerm]
    splits var@(Var (ZVar _ _ ZVariable)) = return [var]
    splits term = do
      maybe_expr <- liftZenoReader $ applyFunction [] term
      case maybe_expr of
        Nothing -> return []
        Just (Value term') -> return [term']
        Just (CaseOf term' _) -> do
          terms <- splits term'
          return (term : terms)      
 
invalidTerm :: ZTerm -> Bool
invalidTerm term = 
  let ((Var func) : args) = flattenTerm term
      arg_types = flattenFunType (zVarType func)
      res_types = drop (length args) arg_types 
  in null res_types
       
randomSequence :: [Solver a] -> Solver [a]
randomSequence solvers = do
  random <- sRandom <$> ask
  let (solvers', random') = shuffle solvers random
  localRandom random' (splitGens solvers')
  where
    localRandom :: StdGen -> Solver a -> Solver a
    localRandom gen = local (\s -> s { sRandom = gen })
        
    splitGens :: [Solver a] -> Solver [a]
    splitGens = foldl' splitGen (return [])
      where 
        splitGen :: Solver [a] -> Solver a -> Solver [a]
        splitGen ss s = do
          random <- sRandom <$> ask
          let (r1, r2) = split random
          ss' <- localRandom r2 ss
          s' <- localRandom r1 s
          return (s' : ss')
          
substituteTerms :: ZTermSubstitution -> Solver a -> Solver a
substituteTerms sub = 
  let f :: WithinTraversable ZTerm (t ZVar) => t ZVar -> t ZVar
      f = substituteWithin sub 
  in local $ \s -> s
      { sGoal = f (sGoal s), 
        sHypotheses = map f (sHypotheses s) }
        
enableGeneralisation :: Solver a -> Solver a
enableGeneralisation = local $ \s -> s { sDoGeneralisation = True }

disableGeneralisation :: Solver a -> Solver a
disableGeneralisation = local $ \s -> s { sDoGeneralisation = False }
    
localGoal :: ZClause -> Solver a -> Solver a
localGoal goal = local $ \s -> s { sGoal = goal }

localGoalEquality :: ZEquality -> Solver a -> Solver a
localGoalEquality eq = local $ \s -> s 
  { sGoal = (sGoal s) { clauseEquality = eq } }

localGoalConditions :: [ZEquality] -> Solver a -> Solver a
localGoalConditions conds = local $ \s -> s
  { sGoal = (sGoal s) { clauseConditions = conds } }
  
addCond :: ZEquality -> Solver a -> Solver a
addCond cond = local $ \s -> s
  { sGoal = (sGoal s) 
    { clauseConditions = cond : (clauseConditions (sGoal s)) } }

addConditions :: [ZEquality] -> Solver a -> Solver a
addConditions = foldl (\f c -> f . addCond c) id
    
addHypothesis :: ZClause -> Solver a -> Solver a
addHypothesis cls = local $ \s -> s
  { sHypotheses = cls : (sHypotheses s) }
  
removeHypothesis :: ZClause -> Solver a -> Solver a
removeHypothesis cls = local $ \s -> s
  { sHypotheses = delete cls (sHypotheses s) }
  
removeHypotheses :: [ZClause] -> Solver a -> Solver a
removeHypotheses = foldl (\f c -> f . removeHypothesis c) id

removeAllHypotheses :: Solver a -> Solver a
removeAllHypotheses = local $ \s -> s { sHypotheses = [] }

addHypotheses :: [ZClause] -> Solver a -> Solver a
addHypotheses = foldl (\f c -> f . addHypothesis c) id
  
debugShowGoal :: Solver a -> Solver a
debugShowGoal solver = do
  goal <- sGoal <$> ask
  shown <- liftZenoReader $ showZClause goal
  traceZ shown $ solver
  
generalisation :: Solver [ZProof] -> Solver [ZProof]
generalisation step = do
  do_it <- sDoGeneralisation <$> ask
  if do_it then step else return [Unknown]

takeIndStep :: Solver ZProof -> Solver ZProof
takeIndStep = local $ \s -> s { sIndSteps = (sIndSteps s) - 1 }
      
pureRecordType :: ZType -> Bool
pureRecordType z_type = 
  let ((TypeVar dtype@(ZDataType {})) : _) = flattenAppType z_type
  in  not (recursiveTypeVar dtype)
      
doIndStep :: Solver [ZProof] -> Solver [ZProof]
doIndStep solver = do
  inds <- sIndSteps <$> ask
  if inds <= 0
    then return []
    else solver
            
proofstep :: String -> [ZClause] -> Solver ZProofStep
proofstep name added = do
  goal <- fullGoal
  return (PL.ProofStep name (Just goal) added) 
